import torch
import torch.nn as nn
import torch.nn.functional as F
from nets.custom_loss import my_loss
from layers.gcn_layer import GCNLayer
from layers.mlp_readout_layer import MLPReadout

"""
    GCN: Graph Convolutional Networks
    Thomas N. Kipf, Max Welling, Semi-Supervised Classification with Graph Convolutional Networks (ICLR 2017)
    http://arxiv.org/abs/1609.02907
"""


class GCNNet(nn.Module):

    def __init__(self, net_params):
        super().__init__()

        in_dim = net_params.in_dim_node
        in_dim_loop = net_params.in_dim_loop
        hidden_dim = net_params.hidden_dim
        out_dim = net_params.out_dim
        n_classes = net_params.n_classes
        in_feat_dropout = net_params.in_feat_dropout
        dropout = net_params.dropout
        n_layers = net_params.n_layers
        self.readout = net_params.readout
        self.batch_norm = net_params.batch_norm
        self.residual = net_params.residual
        self.n_classes = n_classes
        self.device = net_params.device
        self.pos_enc = net_params.pos_enc
        self.pos_enc_dim = net_params.pos_enc_dim
        self.add_loop_type = net_params.add_loop_type
        self.loss_type = net_params.loss_type

        self.embedding_h = nn.Embedding(in_dim, hidden_dim)
        self.embedding_loop = nn.Embedding(in_dim_loop, hidden_dim)
        self.in_feat_dropout = nn.Dropout(in_feat_dropout)

        if self.pos_enc == True:
            self.embedding_pos_enc = nn.Linear(self.pos_enc_dim, hidden_dim)

        self.layers = nn.ModuleList([GCNLayer(hidden_dim, hidden_dim, F.relu, dropout,
                                              self.batch_norm, self.residual) for _ in range(n_layers - 1)])
        self.layers.append(GCNLayer(hidden_dim, out_dim, F.relu, dropout, self.batch_norm, self.residual))
        self.MLP_layer = MLPReadout(out_dim, n_classes)

    def forward(self, g, h, e, loop_type=None, h_pos_enc=None):
        # 嵌入层
        h = self.embedding_h(h)
        # dropout层
        h = self.in_feat_dropout(h)

        if self.add_loop_type:  # 是否使用loop_type特征
            loop_type = self.embedding_loop(loop_type)
            h = h + loop_type

        if self.pos_enc:  # 是否使用位置编码
            h_pos_enc = self.embedding_pos_enc(h_pos_enc.float())
            h = h + h_pos_enc

        # GCN层
        for conv in self.layers:
            h = conv(g, h)
        # output输出层 全连接层
        h_out = self.MLP_layer(h)

        return h_out

    def loss(self, pred, label, mask):
        loss = my_loss(pred, label, mask, self.loss_type)
        return loss

    # def loss(self, pred, label):
    #     # 计算损失值，回归问题，使用MSE损失
    #     # reduce:
    #     # (1)如果 reduce = False，那么 size_average 参数失效，直接返回向量形式的 loss
    #     # (2)如果 reduce = True，那么 loss 返回的是标量
    #     # size_average:
    #     # (1)如果 size_average = False，返回 loss.sum()
    #     # (2)如果 size_average = True，返回 loss.mean()
    #     # criterion = nn.MSELoss(reduce=True, size_average=True)
    #     pred = pred.to(torch.float32)
    #     label = label.to(torch.float32)
    #     # criterion = nn.MSELoss()
    #     # loss = criterion(pred, label)
    #     mcr_mse_loss = MCRMSELoss()
    #     loss = mcr_mse_loss(pred, label)
    #     return loss
